from typing import Literal, TypedDict, List, Any, Dict, Callable
import networkx as nx
import sys
from .knowledge_graph import NetworkXKG
if sys.version_info < (3, 11):
    from typing_extensions import TypedDict, Required, NotRequired
else:
    from typing import TypedDict, Required, NotRequired


# 定义节点和边的结构
class Node(TypedDict):
    id: int
    label: str
    group: str
    value: int  # 假设value是int类型，根据实际情况调整
    data: Any
    x: NotRequired[int]
    y: NotRequired[int]
    fixed: NotRequired[bool]


class Edge(TypedDict):
    from_: int  # 使用from_以避免Python关键字冲突
    to: int
    length: int
    width: int
    label: str
    color: str
    fontColor: str


# 提取节点数据
def extract_nodes(
    graph: NetworkXKG
) -> List[Node]:
    """
    graph: 输入图
    label_formatter: 格式化节点标签的函数，输入为节点的标识以及节点的数据，输出标签

    """

    nodes: List[Node] = []
    for node_id, data in graph.nodes(data=True):
        node: Node = {
            "id": data.get("id", node_id),
            "label": graph.format_node_label(node_id, data),
            "group": data.get("type", "default"),
            "value": data.get("value", 1),  # 默认值1，根据实际情况调整
            "data": data,
        }
        if "x" in data:
            node["x"] = data["x"]
        if "y" in data:
            node["y"] = data["y"]
        if "fixed" in data:
            node["fixed"] = True
        nodes.append(node)
    return nodes


def extract_edges(
    graph: NetworkXKG
) -> List[Edge]:

    edges: List[Edge] = []
    for u, v, data in graph.edges(data=True):
        edge: Edge = {
            "from": u,
            "to": v,
            "length": 50,
            "width": data.get("width", 1),
            "label": graph.format_edge_label((u, v), data),
            "color": data.get("color", "gray"),
            "fontColor": data.get("fontColor", "gray"),
            "data": data,
        }
        edges.append(edge)
    return edges
